'''
Functions for calculating complexity of numeric types.
'''
import sympy as sp
from ..primes import Primes
from .base import complexity

#===============================================================================
# Integers
#===============================================================================
#
# Integers are represented internally as a list of exponents of prime factors.
# All internal functions receive the prefix "i", marking that it works with
# this list of factors rather than the numbers themselves.
#
# Given a list of factors L, the number is constructed as 
#             (-1)**L[0] * 2**L[1] * 3**L[2] * 5**L[3] * ...
# with all other prime factors following. Zero is represented by an empty list
#

# Internal list of prime numbers
PRIMES = Primes([-1, 2])

# Complexity calculation -------------------------------------------------------
@complexity.register(int, sp.Integer)
def int_complexity(N, method='default'):
    r'''Compute the complexity of the given integer N
    
    Examples
    --------
    
    Complexities of numbers 0 to 10
    
    >>> print('Numbers:', list(range(1, 11))); \
    ... print('  Mixed:', [ complexity(i) for i in range(1, 11) ]); \
    ... print('    Mul:', [ complexity(i, 'mul') for i in range(1, 11) ]); \
    ... print('    Add:', [ complexity(i, 'add') for i in range(1, 11) ])
    Numbers: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
      Mixed: [1, 2, 3, 3, 4, 4, 5, 4, 4, 3]
        Mul: [0, 2, 3, 3, 4, 4, 5, 4, 5, 2]
        Add: [1, 2, 2, 3, 4, 4, 4, 4, 4, 4]
    '''

    if method == 'mixed' or method == 'default':
        return int_mixed_complexity(N)
    elif method == 'mul':
        return int_mul_complexity(N)
    elif method == 'add':
        return int_add_complexity(N)
    else:
        raise ValueError('invalid method for integers: %s' % method)

def int_mul_complexity(N):
    '''
    Compute the multiplication complexity of the number from its list of 
    factors.
    
    The multiplication complexity tries to capture the difficulty of using
    the given number in multiplication and division operations. 
    
    We use a pseudo prime factorization in which -1 and 10 are considered 
    as atomic factors as the regular prime numbers. The number -200, for 
    instance, is factorized as ``-200 = (-1)^1 * 10^2 * 2^1``, rather than using 
    the prime factorization ``-200 = (-1) * 2^3 + 5^2``.
    
    A complexity value is assigned to each factor according to their 
    position in the list of prime numbers. The table shows the complexity
    of the first few factors
    
        +------------+----+----+----+----+----+----+----+----+----+-----+
        | Factor     |  0 | -1 | 10 |  2 |  3 |  5 |  7 | 11 | 13 | ... |  
        +------------+----+----+----+----+----+----+----+----+----+-----+
        | Complexity |  1 |  1 |  2 |  2 |  3 |  4 |  5 |  6 |  7 | ... |
        +------------+----+----+----+----+----+----+----+----+----+-----+
    
    Given the factorization, the multiplication complexity of the number is
    computed as the sum of <factor complexity> * <exponent> of each term in 
    the factorization minus the number of multiplication signs in the full
    factorization.
    
    Using the number ``-200`` as an example, we have a complexity of:: 
    
        C(-200) = C(-1)*1 + C(10)*2 + C(2)*1 - 2
                =  1 * 1  +  3 * 2  +  2 * 1 - 2
                =  7
                
    The multiplication complexity of 0 and 1 is 0.0.
    '''

    # Special cases 0, 1, -1
    if N in (0, 1, -1):
        return 0.0

    # Compute raw complexity
    factors = sp.factorint(abs(N))

    # power 10 factorization
    try:
        e2 = factors[2]
        e5 = factors[5]
        e10 = min(e2, e5)
        factors[2] -= e10
        factors[5] -= e10
    except KeyError:
        e10 = 0

    # Compute raw factorization
    prime_index = PRIMES.index
    compl = sum((prime_index(p) + 1) * e for (p, e) in factors.items())
    # Add contribution of base 10 factorization
    compl += 2 * e10
    # Extra complexity for negative numbers
    compl += 0.0 if N >= 0 else 1.0
    # Reduce complexity penalty of long numbers
    compl -= e10 + sum(factors.values()) - 1
    return float(compl)

def int_add_complexity(N):
    '''
    Compute the addition complexity of the number.
    
    The addition complexity tries to capture the difficulty of using
    the given number in addition and subtraction operations.
    
    List of complexities for the first 10 numbers 
    >>> for i in range(1, 11):
    ...     print('%2s: %s' % (i, complexity(i, 'add', round=False)))
     1: 1.0
     2: 2.0
     3: 2.5
     4: 3.0
     5: 3.5
     6: 4.0
     7: 4.5
     8: 4.5
     9: 3.5
    10: 3.5
    '''

    # Special cases 0, 1, -1
    if N == 0:
        return 0.0
    elif N == 1 or N == -1:
        return 1.0

    # Compute complexity
    C = {'0': 1.0, '1': 1.5, '2': 2.0, '3': 2.5, '4': 3.0, '5': 3.5,
         '6': 4.0, '7': 4.5, '8': 4.5, '9': 3.5, }
    digits = str(abs(N))
    num_zeros = digits.count('0')

    # Add complexity of each digit
    compl = sum(C[x] for x in digits)
    # Add complexity for long numbers ...
    compl += (len(digits) - 1)
    # ... but bound it for numbers with lots of zeros
    compl -= max(num_zeros - 3, 0)
    # Negative numbers also have a complexity bonus 
    if N < 0:
        compl += 1
    return compl

def int_mixed_complexity(N):
    '''Compute the combined add/mul complexity of the number. The combined
    complexity is the weighted average between the largest complexity (with 
    weight 2/3) and the smallest one (with weight 1/3).
    '''

    m, a = int_add_complexity(N), int_mul_complexity(N)
    A = max(m, a)
    B = (m + a) / 2
    return (A + 2 * B) / 3

#===============================================================================
# Rational
#===============================================================================
@complexity.register(sp.Rational)
def rat_complexity(N, method='default'):
    '''
    Compute the complexity of a rational number.
    '''
    a, b = int(sp.numer(N)), int(sp.denom(N))
    C_A = int_complexity(a, method)
    C_B = int_complexity(b, method)
    C_M = max(a, b)
    return 0.75 * (C_A + C_B) + 0.5 * C_M

#===============================================================================
# Tests
#===============================================================================
if __name__ == '__main__':
    import doctest
    doctest.testmod()

